#include <sds/sds.h>

#include <cluster.h>
#include <matrix.h>
#include <segment.h>

#include <random.h>
#include <chi_square.h>

#include <vector>
#include <iostream>
#include <iomanip>

using namespace sds;
using namespace std;

Random rng;

#if !defined(CLASS_MAXLH) && !defined(CLASS_BHAT)
#  define CLASS_BHAT
#endif

// Parameters

double noise = 0.0001;

//double leave = 0.00001;

double min_relevant = 0.25;
double min_rate = 0.25;

//double scaleRate = 1.1;

// Constans

//double leaveChiSquare = invChiSquare(1.0 - leave, CH);


int main(int argc, const char* argv[]) {
  SFile segfile(argv[2]);
  SDataView<unsigned char> seg =
    segfile.getDataset(0).getDataView<unsigned char>();

  SFile reffile(argv[3]);
  SDataView<unsigned char> ref =
    reffile.getDataset(0).getDataView<unsigned char>();

  vector<SInteger> size = ref.getSize();

  SFile srcfile(argv[1]);
  SDataView<Point> src(size);

  for (int k = 0; k < CH; ++k) {
    SDataView<unsigned char> view =
      srcfile.getDataset(k).getDataView<unsigned char>();
    for (int i = 0; i < size[0]; ++i) {
      for (int j = 0; j < size[1]; ++j) {
	((Point&)src[i][j])[k] = (double)view[i][j];
      }    
    }
  }
  
  int num = 0;
  Cluster signs[256];
  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      if (ref[i][j] != 0) {
	signs[ref[i][j]] += src[i][j];
	++num;
      }
    }
  }

  vector<unsigned char> index;
  vector<Vector> expected;
  vector<Matrix> inverse;
  vector<double> lndet;

  for (int k = 1; k < 256; ++k) {
    if (signs[k].size() == 0) continue;
    index.push_back(k);
    Vector exp(CH);
    for (int i = 0; i < CH; ++i) {
      exp(i) = signs[k].expected(i);
    }
    expected.push_back(exp);
    Matrix cov(CH, CH);
    for (int i = 0; i < CH; ++i) {
      for (int j = 0; j < CH; ++j) {
	cov(i, j) = signs[k].covariance(i, j);
      }
    }
    if (cov.determinant() == 0) {
      for (int i = 0; i < CH; ++i) {
	cov(i, i) += noise;
      }
    }
    lndet.push_back(log(fabs(cov.determinant())));
    inverse.push_back(cov.inverse());
  }

  SDataView<int> ren(size);
  vector<Cluster> ren_signs;
  ReadSegments(src, seg, ren, ren_signs);

  std::vector<int> ren_cla(ren_signs.size());

  for (int s = 0; s < (int)ren_signs.size(); ++s) {
    std::cerr << s << " / " << ren_signs.size() << std::endl;
    Vector seg_exp(CH);
    for (int i = 0; i < CH; ++i) {
      seg_exp(i) = ren_signs[s].expected(i);
    }
#if defined(CLASS_MAXLH)
    Matrix seg_mul(CH, CH);
    for (int i = 0; i < CH; ++i) {
      for (int j = 0; j < CH; ++j) {
        seg_mul(i, j) = ren_signs[s].covariance(i, j) 
          + seg_exp(i) * seg_exp(j);
      }
    }
#endif
    
    int minPos = -1; 
    double minVal = 0.0;
    for (int k = 0; k < (int)index.size(); ++k) {
#if defined(CLASS_MAXLH)
      double dist = mulTrace(inverse[k], seg_mul)
	- 2 * expected[k] * inverse[k] * seg_exp 
	+ expected[k] * inverse[k] * expected[k];
      dist += lndet[k];
#elif defined(CLASS_BHAT)
      Matrix avg_cov(CH, CH);
      for (int i = 0; i < CH; ++i) {
	for (int j = 0; j < CH; ++j) {
	  avg_cov(i, j) =
	    (signs[index[k]].covariance(i, j) + 
	     ren_signs[s].covariance(i, j)) / 2;
	}
      }
      if (avg_cov.determinant() == 0.0) {
	for (int i = 0; i < CH; ++i) {
	  avg_cov(i, i) += noise;
	}
      }
      Vector exp_dist(CH);
      for (int i = 0; i < CH; ++i) {
	exp_dist(i) = signs[index[k]].expected(i) - ren_signs[s].expected(i);
      }
      double dist = 0.125 * exp_dist * avg_cov.inverse() * exp_dist +
	0.5 * (log(fabs(avg_cov.determinant())) - 0.5 * lndet[k]);
#endif
      if (minPos == -1 || minVal > dist) {
        minPos = k;
        minVal = dist;
      }
    }
    ren_cla[s] = minPos != -1 ? index[minPos] : 0;
  }

  SDataView<unsigned char> trg(size);
  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      trg[i][j] = ren_cla[ren[i][j]];
    }
  }


  SFile trgfile(argv[4], SFile::CREATE);
  trgfile.createDataset("clasified", UINT8, size).
    setDataView<unsigned char>(trg);
  return 0;
}
